import gradio as gr
import httpx
import html
import logging
import re
from datetime import datetime, timedelta
import os
import json

# =========================
# 日志设置
# =========================
def setup_logging():
    os.makedirs("logs", exist_ok=True)
    log_filename = f"logs/gradio_ui_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler(log_filename, encoding="utf-8"),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)

logger = setup_logging()

# =========================
# API 配置
# =========================
API_BASE_URL = "http://localhost:12344"
API_ENDPOINTS = {
    "general": f"{API_BASE_URL}/chat/rag",
    "rag": f"{API_BASE_URL}/chat/rag",
    "weather": f"{API_BASE_URL}/chat/weather",
    "flight": f"{API_BASE_URL}/chat/flight",
    "ticket": f"{API_BASE_URL}/chat/ticket",
    "map": f"{API_BASE_URL}/chat/map"
}

# =========================
# CSS 样式
# =========================
css = """/* 基础样式 */
.main-container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}

.header-text {
    text-align: center;
    margin-bottom: 20px;
}

.service-tabs {
    margin-bottom: 20px;
}

/* 消息输入容器 */
.message-input-container {
    margin-top: 20px;
    gap: 10px;
}

/* 按钮样式 */
.custom-button {
    padding: 8px 16px;
    border: none;
    border-radius: 4px;
    cursor: pointer;
}

.custom-button:not(.secondary) {
    background: #1890ff;
    color: white;
}

.custom-button.secondary {
    background: #f5f5f5;
    color: #666;
    border: 1px solid #ddd;
}

/* 聊天区域 */
#general_html, #rag_html, #weather_html, 
#flight_html, #ticket_html, #map_html {
    min-height: 400px;
    max-height: 600px;
    overflow-y: auto;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 16px;
    background: #fafafa;
}

/* 用户消息 */
.user-message {
    margin-bottom: 16px;
    display: flex;
    justify-content: flex-end;
}

.user-message .answer-block {
    background: #f0f0f0; /* 改为灰色背景 */
    color: #333; /* 改为深色文字 */
    padding: 12px 16px;
    border-radius: 12px 12px 0 12px;
    max-width: 80%;
}

/* 机器人消息 - 垂直排列 */
.bot-message {
    margin-bottom: 16px;
    display: flex;
    flex-direction: column;
    gap: 8px;
}

/* 思考区块 */
.thinking-block {
    background: #f8f9fa;
    border: 1px solid #e9ecef;
    border-radius: 8px;
    padding: 12px;
    position: relative;
}

.thinking-block::before {
    content: "💭 思考过程";
    display: block;
    font-weight: bold;
    font-size: 12px;
    color: #6c757d;
    margin-bottom: 8px;
    border-bottom: 1px dashed #dee2e6;
    padding-bottom: 4px;
}

.thinking-content {
    white-space: pre-wrap;
    font-size: 14px;
    line-height: 1.4;
    color: #495057;
}

/* 答案区块 */
.answer-block {
    background: white;
    border: 1px solid #dee2e6;
    border-radius: 8px;
    padding: 12px;
    position: relative;
}

.answer-block::before {
    content: "💡 正式回答";
    display: block;
    font-weight: bold;
    font-size: 12px;
    color: #000; /* 改为黑色 */
    margin-bottom: 8px;
    border-bottom: 1px dashed #dee2e6;
    padding-bottom: 4px;
}

.answer-content {
    white-space: pre-wrap;
    font-size: 15px;
    line-height: 1.5;
    color: #212529;
}

/* JSON格式指南 */
.json-guide {
    background: #e7f3ff;
    border-left: 4px solid #1890ff;
    padding: 12px;
    border-radius: 4px;
    margin-bottom: 16px;
}

/* 标签页样式 */
.tabs {
    border-bottom: 2px solid #e0e0e0;
}

.tab-item {
    padding: 12px 24px;
    cursor: pointer;
    border-bottom: 3px solid transparent;
}

.tab-item.selected {
    border-bottom-color: #1890ff;
    color: #1890ff;
    font-weight: bold;
}

/* 表单样式 */
.ticket-form {
    background: #f8f9fa;
    padding: 20px;
    border-radius: 8px;
    border: 1px solid #e9ecef;
}

.form-group {
    margin-bottom: 16px;
}

.form-label {
    font-weight: bold;
    margin-bottom: 8px;
    display: block;
    color: #495057;
}

/* 滚动条样式 */
::-webkit-scrollbar {
    width: 6px;
}

::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 3px;
}

::-webkit-scrollbar-thumb {
    background: #c1c1c1;
    border-radius: 3px;
}

::-webkit-scrollbar-thumb:hover {
    background: #a8a8a8;
}"""

# =========================
# 流式响应函数
# =========================
async def call_api_stream_charwise(endpoint: str, payload: dict):
    buffer = bytearray()
    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", endpoint, json=payload) as response:
            async for chunk in response.aiter_bytes():
                buffer.extend(chunk)
                while True:
                    try:
                        decoded = buffer.decode("utf-8")
                        for char in decoded:
                            yield char
                        buffer.clear()
                        break
                    except UnicodeDecodeError:
                        break
            if buffer:
                decoded = buffer.decode("utf-8", errors="replace")
                for char in decoded:
                    yield char

# =========================
# 内容解析与格式化
# =========================
def parse_multiple_thinking_blocks(content: str):
    think_pattern = r'<think>(.*?)</think>'
    thinking_blocks = re.findall(think_pattern, content, re.DOTALL)
    answer_content = re.sub(think_pattern, '', content, flags=re.DOTALL).strip()
    return [b.strip() for b in thinking_blocks if b.strip()], answer_content

def format_assistant_message(content: str):
    thinking_blocks, answer_content = parse_multiple_thinking_blocks(content)
    blocks = []
    for block in thinking_blocks:
        blocks.append(f"""
        <div class="thinking-block">
            <div class="thinking-content">{html.escape(block)}</div>
        </div>
        """)
    if answer_content:
        blocks.append(f"""
        <div class="answer-block">
            <div class="answer-content">{html.escape(answer_content)}</div>
        </div>
        """)
    if not blocks:
        blocks.append(f"""
        <div class="answer-block">
            <div class="answer-content">{html.escape(content)}</div>
        </div>
        """)
    return "".join(blocks)

# =========================
# 核心响应逻辑（HTML版本）
# =========================
def make_respond_fn(service_type):
    async def respond_fn(message, chat_html):
        if not message or not message.strip():
            yield "", chat_html
            return

        user_block = f"""
        <div class="user-message">
            <div class="answer-block">
                <div class="answer-content">{html.escape(message)}</div>
            </div>
        </div>
        """
        chat_html += user_block
        yield "", chat_html

        endpoint = API_ENDPOINTS.get(service_type)
        payload = {"message": message} if service_type != "ticket" else message

        assistant_message = ""
        assistant_block = ""

        async for char in call_api_stream_charwise(endpoint, payload):
            assistant_message += char
            formatted = format_assistant_message(assistant_message)
            assistant_block = f"""
            <div class="bot-message">{formatted}</div>
            """
            yield "", chat_html + assistant_block

        chat_html += assistant_block
        yield "", chat_html

    return respond_fn

def clear_history():
    return "", ""

def generate_ticket_query(from_station, to_station, train_date):
    """生成火车票查询JSON"""
    if not from_station or not to_station or not train_date:
        raise gr.Error("请填写完整的查询信息！")
    
    query_data = {
        "from_station": from_station.strip(),
        "to_station": to_station.strip(),
        "train_date": train_date.strip()
    }
    return json.dumps(query_data, ensure_ascii=False)

def generate_date_options():
    """生成今天到15天内的日期选项"""
    today = datetime.now()
    dates = []
    for i in range(15):  # 0到14天，共15天
        date = today + timedelta(days=i)
        date_str = date.strftime("%Y-%m-%d")
        display_str = date.strftime("%Y年%m月%d日")
        if i == 0:
            display_str += " (今天)"
        elif i == 1:
            display_str += " (明天)"
        elif i == 2:
            display_str += " (后天)"
        dates.append((display_str, date_str))
    return dates

# =========================
# UI 构建
# =========================
def create_travel_assistant_ui():
    with gr.Blocks(title="旅行助手", css=css, theme=gr.themes.Soft()) as demo:
        with gr.Column(elem_classes=["main-container"]):
            gr.Markdown("# 🚀 旅行助手 - Travel Assistant", elem_classes=["header-text"])

            with gr.Tabs(elem_classes=["service-tabs"]):
                # 💬 通用对话
                with gr.TabItem("💬 通用对话"):
                    general_html = gr.HTML(value="", elem_id="general_html")
                    with gr.Row(elem_classes=["message-input-container"]):
                        general_msg = gr.Textbox(placeholder="请输入您的问题...", lines=2, max_lines=4)
                        general_submit = gr.Button("发送", elem_classes=["custom-button"])
                        general_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"])

                # 📋 差旅政策
                with gr.TabItem("📋 差旅政策"):
                    rag_html = gr.HTML(value="", elem_id="rag_html")
                    with gr.Row(elem_classes=["message-input-container"]):
                        rag_msg = gr.Textbox(placeholder="请输入关于差旅政策的问题，例如：关于出差的住宿费规定？", lines=2, max_lines=4)
                        rag_submit = gr.Button("发送", elem_classes=["custom-button"])
                        rag_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"])

                # 🌤️ 天气查询
                with gr.TabItem("🌤️ 天气查询"):
                    weather_html = gr.HTML(value="", elem_id="weather_html")
                    with gr.Row(elem_classes=["message-input-container"]):
                        weather_msg = gr.Textbox(placeholder="例如：重庆明天天气如何？", lines=2, max_lines=4)
                        weather_submit = gr.Button("发送", elem_classes=["custom-button"])
                        weather_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"])

                # ✈️ 航班查询
                with gr.TabItem("✈️ 航班查询"):
                    flight_html = gr.HTML(value="", elem_id="flight_html")
                    with gr.Row(elem_classes=["message-input-container"]):
                        flight_msg = gr.Textbox(placeholder="例如：我要买明天北京飞宁波的航班，几点出发无所谓，不过最好提供餐食", lines=2, max_lines=4)
                        flight_submit = gr.Button("发送", elem_classes=["custom-button"])
                        flight_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"])

                # 🚄 火车票查询
                with gr.TabItem("🚄 火车票查询"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("### 🎫 填写查询信息")
                            with gr.Group(elem_classes=["ticket-form"]):
                                with gr.Column():
                                    from_station = gr.Textbox(
                                        label="出发站",
                                        placeholder="例如：宁波",
                                        elem_classes=["form-group"]
                                    )
                                    to_station = gr.Textbox(
                                        label="到达站", 
                                        placeholder="例如：杭州",
                                        elem_classes=["form-group"]
                                    )
                                    # 使用下拉选择器作为日期选择
                                    train_date = gr.Dropdown(
                                        label="出发日期",
                                        choices=generate_date_options(),
                                        value=datetime.now().strftime("%Y-%m-%d"),
                                        elem_classes=["form-group"]
                                    )
                                query_btn = gr.Button("🚄 查询车票", elem_classes=["custom-button"], size="lg")
                            
                        with gr.Column(scale=2):
                            gr.Markdown("### 💬 查询结果")
                            ticket_html = gr.HTML(value="", elem_id="ticket_html")
                    
                    # 隐藏的文本框，用于存储生成的JSON
                    ticket_msg = gr.Textbox(visible=False)
                    ticket_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"], visible=False)

                # 🗺️ 地图服务
                with gr.TabItem("🗺️ 地图服务"):
                    map_html = gr.HTML(value="", elem_id="map_html")
                    with gr.Row(elem_classes=["message-input-container"]):
                        map_msg = gr.Textbox(placeholder="例如：给我一份北京大学到天安门的路径规划？", lines=2, max_lines=4)
                        map_submit = gr.Button("发送", elem_classes=["custom-button"])
                        map_clear = gr.Button("清空", elem_classes=["custom-button", "secondary"])

            gr.Markdown("*© 2025 旅行助手 - 为您的旅程提供全方位服务*")

        # 绑定事件
        for (msg, html_box, submit, clear, t) in [
            (general_msg, general_html, general_submit, general_clear, "general"),
            (rag_msg, rag_html, rag_submit, rag_clear, "rag"),
            (weather_msg, weather_html, weather_submit, weather_clear, "weather"),
            (flight_msg, flight_html, flight_submit, flight_clear, "flight"),
            (map_msg, map_html, map_submit, map_clear, "map")
        ]:
            msg.submit(make_respond_fn(t), inputs=[msg, html_box], outputs=[msg, html_box])
            submit.click(make_respond_fn(t), inputs=[msg, html_box], outputs=[msg, html_box])
            clear.click(clear_history, outputs=[html_box, msg])

        # 火车票查询特殊处理
        query_btn.click(
            fn=generate_ticket_query,
            inputs=[from_station, to_station, train_date],
            outputs=ticket_msg
        ).then(
            fn=make_respond_fn("ticket"),
            inputs=[ticket_msg, ticket_html],
            outputs=[ticket_msg, ticket_html]
        )

    return demo

if __name__ == "__main__":
    demo = create_travel_assistant_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=True)